#!/bin/bash

# ==============================================
# RDP Proxy 安装脚本 (静态编译版) - 修复版
# 版本: 2.3
# 更新日期: 2024-05-21
# ==============================================

# 配置项
readonly DOWNLOAD_URL="https://gitee.com/zuoz_net/windows/raw/master/rdp-proxy-2"
readonly INSTALL_DIR="/opt/rdp-proxy"
readonly SERVICE_NAME="rdp-proxy"
readonly LISTEN_PORT="3389"

# 颜色定义
readonly CR="\033[0m"
readonly RED="\033[31m"
readonly GREEN="\033[32m"
readonly YELLOW="\033[33m"
readonly BLUE="\033[34m"

print_msg() {
    local color=$1
    shift
    echo -e "${color}$*${CR}" >&2
}

check_root() {
    [ "$(id -u)" -eq 0 ] || {
        print_msg $RED "✗ 必须使用root用户运行此脚本！"
        exit 1
    }
}

cleanup_old() {
    print_msg $BLUE "[1/6] 清理旧版本..."
    systemctl stop "$SERVICE_NAME" 2>/dev/null || true
    systemctl disable "$SERVICE_NAME" 2>/dev/null || true
    pkill -9 -f "$INSTALL_DIR/rdp-proxy" 2>/dev/null || true
    rm -rf "$INSTALL_DIR" "/etc/systemd/system/${SERVICE_NAME}.service"
    sleep 2
}

install_deps() {
    print_msg $BLUE "[2/6] 安装系统依赖..."
    if ! command -v wget &>/dev/null && ! command -v curl &>/dev/null; then
        if command -v apt-get &>/dev/null; then
            apt-get update && apt-get install -y wget curl
        elif command -v yum &>/dev/null; then
            yum install -y wget curl
        else
            print_msg $YELLOW "⚠️ 无法自动安装wget/curl，请手动安装"
            exit 1
        fi
    fi
}

download_binary() {
    print_msg $BLUE "[3/6] 下载程序文件..."
    local temp_file
    temp_file=$(mktemp "/tmp/rdp-proxy.XXXXXX") || {
        print_msg $RED "✗ 无法创建临时文件"
        exit 1
    }

    # 尝试wget下载
    if command -v wget &>/dev/null; then
        print_msg $YELLOW "尝试使用wget下载..."
        if wget --quiet --tries=3 --timeout=15 -O "$temp_file" "$DOWNLOAD_URL" && [ -s "$temp_file" ]; then
            chmod +x "$temp_file"
            echo "$temp_file"
            return 0
        fi
    fi

    # 尝试curl下载
    if command -v curl &>/dev/null; then
        print_msg $YELLOW "wget失败，尝试使用curl..."
        if curl -sSL -o "$temp_file" "$DOWNLOAD_URL" && [ -s "$temp_file" ]; then
            chmod +x "$temp_file"
            echo "$temp_file"
            return 0
        fi
    fi

    print_msg $RED "✗ 下载失败，可能原因:"
    print_msg $YELLOW "1. URL不可访问: $DOWNLOAD_URL"
    print_msg $YELLOW "2. 网络连接问题"
    print_msg $YELLOW "3. 服务器限制"
    print_msg $GREEN "✅ 解决方案:"
    print_msg $GREEN "请手动下载文件并放置到 $INSTALL_DIR/rdp-proxy"
    print_msg $GREEN "然后执行: chmod +x $INSTALL_DIR/rdp-proxy"
    rm -f "$temp_file"
    exit 1
}

install_service() {
    local src_file="$1"
    [ -f "$src_file" ] || {
        print_msg $RED "✗ 源文件不存在: $src_file"
        exit 1
    }

    print_msg $BLUE "[4/6] 安装服务文件..."
    mkdir -p "$INSTALL_DIR" || {
        print_msg $RED "✗ 无法创建目录 $INSTALL_DIR"
        exit 1
    }

    mv -f "$src_file" "$INSTALL_DIR/rdp-proxy" || {
        print_msg $RED "✗ 文件移动失败"
        exit 1
    }

    cat > "/etc/systemd/system/${SERVICE_NAME}.service" <<EOF
[Unit]
Description=RDP Proxy Service
After=network.target

[Service]
Type=simple
ExecStart=$INSTALL_DIR/rdp-proxy
Restart=always
User=root

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload || {
        print_msg $RED "✗ systemd配置重载失败"
        exit 1
    }
}

setup_firewall() {
    print_msg $BLUE "[5/6] 配置防火墙..."
    if command -v ufw &>/dev/null && ufw status | grep -q "active"; then
        ufw allow "$LISTEN_PORT/tcp"
    elif command -v firewall-cmd &>/dev/null; then
        firewall-cmd --add-port="${LISTEN_PORT}/tcp" --permanent
        firewall-cmd --reload
    else
        print_msg $YELLOW "⚠️ 未检测到活动防火墙，请手动放行端口 $LISTEN_PORT"
    fi
}

start_service() {
    print_msg $BLUE "[6/6] 启动服务..."
    if systemctl enable --now "$SERVICE_NAME"; then
        print_msg $GREEN "✅ 安装成功！"
        echo -e "服务状态: $(systemctl is-active $SERVICE_NAME)"
        echo -e "监听端口: $LISTEN_PORT"
    else
        print_msg $RED "✗ 服务启动失败"
        journalctl -u "$SERVICE_NAME" --no-pager -n 10
        exit 1
    fi
}

main() {
    check_root
    cleanup_old
    install_deps
    local temp_path
    temp_path=$(download_binary)
    install_service "$temp_path"
    setup_firewall
    start_service
}

main "$@"
